import os
import json
import logging
from collections import defaultdict
from pathlib import Path
from code.image_style_trans.style_processor import StyleProcessor

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class ImageProcessor:
    def __init__(self, project_root, output_dir):
        """
        初始化单词数据集处理器
        """
        # 初始化风格处理器
        self.style_processor = StyleProcessor()
        self.project_root = project_root
        self.output_dir = output_dir

        # 检查服务可用性
        if not self.style_processor.available:
            logger.warning("风格处理器初始化失败，图像风格转换功能将不可用")

    def process_word_dataset(self, dataset_path, prompt=None, negative_prompt=None):
        """
        处理单词数据集，对每个单词的句子图像进行风格转换
        :param dataset_path: 单词数据集JSON文件路径
        :param prompt: 风格描述提示词，如果为None则使用默认提示词
        :param negative_prompt: 负向提示词
        :return: 更新后的数据集
        """
        if not self.style_processor.available:
            logger.error("风格处理器未正确初始化，无法处理单词数据集")
            return None

        new_dataset_record = defaultdict(list)
        # 读取默认提示词
        if prompt is None:
            prompt_path = Path(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) / "prompt" / "image_style_trans.txt"
            if prompt_path.exists():
                with open(prompt_path, "r", encoding="utf-8") as f:
                    prompt = f.read().strip()
            else:
                prompt = "转换为蜡笔彩绘风格"

        logger.info(f"使用提示词: {prompt}")

        # 读取数据集
        dataset_path = Path(dataset_path)
        new_dataset_path = dataset_path.parent / f"{dataset_path.stem}_with_style{dataset_path.suffix}"

        if new_dataset_path.exists():
            with open(new_dataset_path, "r", encoding="utf-8") as f:
                dataset = json.load(f)
        else:
            with open(dataset_path, "r", encoding="utf-8") as f:
                dataset = json.load(f)

        try:
            # 获取IP ID
            ip_id = None
            for word, sentences in dataset.items():
                if sentences and len(sentences) > 0:
                    ip_id = sentences[0].get("ip_id")
                    break

            if not ip_id:
                logger.error("无法从数据集中获取IP ID")
                return None

            logger.info(f"处理IP ID: {ip_id}")

            # 创建风格转换图像目录
            image_output_dir = Path(os.path.join(self.output_dir, "images/style_trans")) / ip_id
            os.makedirs(image_output_dir, exist_ok=True)

            # 处理每个单词的句子图像
            total_sentences = sum(len(sentences) for sentences in dataset.values())
            processed_count = 0

            from concurrent.futures import ThreadPoolExecutor
            def process_single_image(sentence, word):
                nonlocal processed_count
                image_id = sentence.get("image_id")
                if 'styled_path' in sentence:
                    processed_count += 1
                    return sentence
                if not image_id:
                    logger.warning(f"句子缺少image_id: {sentence}")
                    processed_count += 1
                    return sentence
                # 设置输出文件名
                output_filename = f"{image_id}.jpg"
                output_path = image_output_dir / output_filename
                if output_path.exists():
                    logger.info(f"图像已存在: {output_path}")
                    styled_url = f"http://159.138.23.13/images/style_trans/{ip_id}/{image_id}.jpg"
                    sentence["styled_path"] = styled_url
                else:
                    # 构建原始图像路径
                    original_image_path = Path(
                        os.path.join(self.output_dir, "images/original")) / ip_id / f"{image_id}.jpg"

                    if not original_image_path.exists():
                        logger.warning(f"原始图像不存在: {original_image_path}")
                        return sentence

                    # 处理图像
                    logger.info(f"处理图像 {processed_count + 1}/{total_sentences}: {original_image_path}")
                    result_path = self.style_processor.process_image(
                        image_path=str(original_image_path),
                        prompt=prompt,
                        negative_prompt=negative_prompt,
                        output_dir=str(image_output_dir),
                        output_filename=output_filename
                    )

                    if result_path:
                        # 更新句子属性
                        styled_url = f"http://159.138.23.13/images/style_trans/{ip_id}/{image_id}.jpg"
                        sentence["styled_path"] = styled_url
                        logger.info(f"图像处理成功: {styled_url}")
                    else:
                        logger.error(f"图像处理失败: {original_image_path}")

                processed_count += 1
                return sentence

            with ThreadPoolExecutor(max_workers=5) as executor:
                words_futures = defaultdict(list)
                for word, sentences in dataset.items():
                    for sentence in sentences[:2]:
                        future = executor.submit(process_single_image, sentence, word)
                        words_futures[word].append(future)

                # 更新数据集
                for word, futures in words_futures.items():
                    for future in futures:
                        new_dataset_record[word].append(future.result())
        except Exception as e:
            logger.error(f"处理单词数据集时发生错误: {str(e)}")

        # 保存更新后的数据集
        with open(new_dataset_path, "w", encoding="utf-8") as f:
            json.dump(new_dataset_record, f, ensure_ascii=False, indent=2)

        logger.info(f"数据集处理完成，成功处理 {processed_count}/{total_sentences} 个图像")
        return dataset


def process_word_dataset(dataset_path, prompt=None, negative_prompt=None, api_key=None):
    """
    处理单词数据集的便捷函数
    :param dataset_path: 单词数据集JSON文件路径
    :param prompt: 风格描述提示词，如果为None则使用默认提示词
    :param negative_prompt: 负向提示词
    :param api_key: 可灵AI API密钥，如果为None则从环境变量获取
    :return: 更新后的数据集
    """
    project_root = Path(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
    output_dir = str(project_root / 'output')
    print(output_dir)
    processor = ImageProcessor(project_root, output_dir)
    return processor.process_word_dataset(dataset_path, prompt, negative_prompt)


if __name__ == "__main__":
    import argparse

    '''
    IP1fef9dba, 疯狂动物城
    IP1750926257, 功夫熊猫1
    IP9754cf1e, 小猪佩奇1
    '''

    dataset_path = "/Users/zhangyuan/projects/auto_video_extract/output/IP1750926257/word_dataset_IP1750926257.json"
    prompt = '去除画面中的文字，动画角色和背景以蜡笔彩绘风格呈现，柔和的色彩过渡搭配粗犷笔触边缘，画面充满手绘质感，保留原始构图基础上添加彩色粉笔画特有的颗粒感，背景融合水彩晕染效果，整体呈现童趣盎然的艺术风格。'

    process_word_dataset(
        dataset_path=dataset_path,
        prompt=prompt,
    )
